import numpy as np
from matplotlib import pyplot as plt
import time
from sklearn.linear_model import LinearRegression
time_now=time.time()
fname='./'+str(int(time_now))[-3:0]
picname = fname+'/'+str(time_now)+'.png'

MU=1
SIGMA=0.1
SAMPLE_NUM=1000
TIME=80
np.random.seed(999)


def plot_wealth(cur_wealth):
    _log_wealth=np.log10(cur_wealth)
    plt.figure()
    plt.hist(_log_wealth, 100, edgecolor='b', facecolor='w')
    plt.xlabel('Wealth(log)')
    plt.ylabel('Number of Individuals')
    plt.xlim((0, 20))

# talent = np.random.normal(MU, SIGMA, SAMPLE_NUM)
# wealth = np.array([10.0] * SAMPLE_NUM)
# luck = np.random.normal(MU, SIGMA, SAMPLE_NUM)
# luck2 = np.random.uniform(0.8, 1.2, SAMPLE_NUM)
# luck_event = np.zeros_like(luck)
# unluck_event = np.zeros_like(luck)
# wealth_wo_luck = np.array([10.0] * SAMPLE_NUM)
# wealth2 = np.array([10.0] * SAMPLE_NUM)
# wealths = []

def simulate():
    # global talent,wealths,luck,luck2,luck_event,unluck_event,wealth_wo_luck,wealth2,wealths
    talent = np.random.normal(MU, SIGMA, SAMPLE_NUM)
    wealth = np.array([10.0] * SAMPLE_NUM)
    luck = np.random.normal(MU, SIGMA, SAMPLE_NUM)
    luck2 = np.random.uniform(0.8, 1.2, SAMPLE_NUM)
    luck_event = np.zeros_like(luck)
    unluck_event = np.zeros_like(luck)
    wealth_wo_luck = np.array([10.0] * SAMPLE_NUM)
    wealth2 = np.array([10.0] * SAMPLE_NUM)
    wealths = []

    for time in range(TIME+1):
        for i in range(SAMPLE_NUM):
            if np.random.rand()*(luck[i])>0.3:
            # if np.random.rand() > 0.3:
                wealth[i]=wealth[i]*(0.4+talent[i]+np.random.normal(0,0.05,1))
                luck_event[i]+=1
            else:
                wealth[i]=wealth[i]*0.5
                unluck_event[i]+=1

            if np.random.rand() > 0.3:
                wealth_wo_luck[i] = wealth_wo_luck[i] * (0.5 + talent[i] + np.random.normal(0, 0.05, 1))
            else:
                wealth_wo_luck[i] = wealth_wo_luck[i] * 0.5

            if np.random.rand()*(luck2[i])>0.3:
                wealth2[i]=wealth[i]*(0.5+talent[i]+np.random.normal(0,0.05,1))
            else:
                wealth2[i]=wealth[i]*0.5

        if time % 10 == 0:
            wealths.append(wealth.copy())
    return talent,wealth,luck,luck2,luck_event,unluck_event,wealth_wo_luck,wealth2,wealths

talent,wealth,luck,luck2,luck_event,unluck_event,wealth_wo_luck,wealth2,wealths=simulate()

def many_people_simu(num):
    talents=[]
    lucks=[]
    for i in range(num):
        talent, wealth, luck, luck2, luck_event, unluck_event, wealth_wo_luck, wealth2, wealths=simulate()
        richest=np.argmax(wealth)
        talents.append(talent[richest])
        lucks.append(luck[richest])
    plt.figure()
    # plt.scatter(talents,lucks)
    # plt.xlabel('Talent')
    # plt.xlim((0.5,1.5))
    # plt.ylim((0.5,1.5))
    # plt.ylabel('Luck')
    # plt.grid

    plt.hist(talents,10,edgecolor='b', facecolor='w')
    plt.xlim((0.5,1.5))
    plt.xlabel('Talent')
    plt.ylabel('Number of Individuals')
many_people_simu(10)

def plot_times():
    for i in range(len(wealths)):
        _log_wealth=np.log10(wealths[i])
        plt.figure()
        plt.title('Year '+str(i*5+1))
        plt.hist(_log_wealth, 100, edgecolor='b', facecolor='w')
        plt.xlabel('Wealth(log10)')
        plt.ylabel('Number of Individuals')
        plt.xlim((0, 18))
        plt.savefig('./npic/'+str(i)+'.png')



# plt.hist(wealth,100,range=(0,1e5),edgecolor='b',facecolor='w')
# plt.xlabel('Wealth')
# plt.ylabel('Number of Individuals')
# plt.ylim((0,100))


#
def plot_all_pics():

    log_wealth = np.log10(wealth)
    log_wealth_wo_luck = np.log10(wealth_wo_luck)
    log_luck = np.log(luck)
    log_wealth2 = np.log10(wealth2)

    plt.figure()
    plt.hist(luck_event,30,edgecolor='b',facecolor='w')
    plt.xlabel('luck_event')
    plt.ylabel('Number')

    plt.figure()
    plt.hist(unluck_event,30,edgecolor='b',facecolor='w')
    plt.xlabel('unluck_event')
    plt.ylabel('Number')

    plt.figure()
    plt.hist(log_wealth,100,edgecolor='b',facecolor='w')
    plt.xlabel('Wealth(log)')
    plt.ylabel('Number of Individuals')
    plt.xlim((0,20))

    plt.figure()
    plt.hist(luck,50,edgecolor='g',facecolor='w')
    plt.xlabel('Luck')
    plt.ylabel('Number of Individuals')

    plt.figure()
    plt.hist(talent,50,edgecolor='r',facecolor='w')
    plt.xlabel('Talent')
    plt.ylabel('Number of Individuals')


    plt.figure()
    plt.scatter(luck,log_wealth,marker='.')
    plt.xlabel('Luck')
    plt.ylabel('Log_wealth')
    # plt.ylim((0,1e6))

    plt.figure()
    plt.scatter(talent,log_wealth,marker='.')
    plt.xlabel('Talent')
    plt.ylabel('Log_wealth')

    plt.figure()
    plt.scatter(talent,log_wealth,marker='.',facecolor='r')
    plt.xlabel('Talent')
    plt.ylabel('Log_wealth')
    # 线性回归
    x = talent.reshape((-1, 1))
    y = log_wealth.reshape(-1, 1)
    l = LinearRegression()
    l.fit(x, y)
    plt.plot([0.6, 1.3], [float(0.6*l.coef_+l.intercept_), float(1.3*l.coef_+l.intercept_)])
    plt.title('R='+str(np.sqrt(l.score(x,y)))[:6])

    plt.figure()
    plt.scatter(talent,log_wealth_wo_luck,marker='.',facecolor='r')
    plt.xlabel('Talent')
    plt.ylabel('Log_wealth_without_luck')
    # 线性回归
    x = talent.reshape((-1, 1))
    y = log_wealth_wo_luck.reshape(-1, 1)
    l = LinearRegression()
    l.fit(x, y)
    plt.plot([0.6, 1.3], [float(0.6 * l.coef_ + l.intercept_), float(1.3 * l.coef_ + l.intercept_)])
    plt.title('R=' + str(np.sqrt(l.score(x, y)))[:6])

    plt.figure()
    plt.scatter(talent,log_wealth2,marker='.')
    plt.xlabel('Talent')
    plt.ylabel('Log_wealth2')

    plt.figure()
    plt.scatter(talent*luck,log_wealth,marker='.',facecolor='r')
    plt.xlabel('Talent * Luck')
    plt.ylabel('Log_wealth')
    # 线性回归
    x = (talent*luck).reshape((-1, 1))
    y = log_wealth.reshape(-1, 1)
    l = LinearRegression()
    l.fit(x, y)
    plt.plot([0.6, 1.5], [float(0.6*l.coef_+l.intercept_), float(1.5*l.coef_+l.intercept_)])
    plt.title('R='+str(np.sqrt(l.score(x,y)))[:6])

    plt.figure()
    plt.scatter(talent,luck,marker='.')
    plt.xlabel('Talent')
    plt.ylabel('Luck')

# plt.savefig(picname)
# plot_all_pics()
# plot_times()
plt.show()